#!/usr/bin/python
import sys, os, re, time
################################################################################
################################################################################
class Params:
   fmt = 'COO'
   verbose = True
   cleanup = True
   xleaf   = False
   trial   = False
   def __init__( self ):
      
      dagMtxHomeEnvName = 'DAGMTX_HOME_DIR'
      
      if not dagMtxHomeEnvName in os.environ:
         print 'Before using ' + sys.argv[0] + ', you need to set the environment variable "' + dagMtxHomeEnvName + '"'
         print 'to point to the home directory for this project.'
         sys.exit( -1 )
      
      self.qtsHomeDir = os.environ[ dagMtxHomeEnvName ]
            
      self.pysHome = os.path.join( self.qtsHomeDir, 'pys' )
      self.binHome = os.path.join( self.qtsHomeDir, 'bin' )
      
      self.hcdDumpBin = os.path.join( self.binHome, 'Mtx2Dag' )
      self.dagReadBin = os.path.join( self.binHome, 'Dag2Mtx' )
      self.mtxDiffBin = os.path.join( self.binHome, 'MtxDiff' )
   

################################################################################
################################################################################
def PrintTimeStamp():
   theTime = time.localtime()
   yr  = '%04d' % theTime.tm_year
   mon = '%02d' % theTime.tm_mon
   day = '%02d' % theTime.tm_mday
   hr  = '%02d' % theTime.tm_hour
   mt  = '%02d' % theTime.tm_min
   sec = '%02d' % theTime.tm_sec
   print '...', '/'.join( [yr, mon, day] ), ':'.join( [hr, mt, sec] )
   sys.stdout.flush()

def SystemWrapperPipelined( cmd, verbose, trial = False ):
   if verbose:
      PrintTimeStamp()
      print '... os.popen(', cmd, ')'
      sys.stdout.flush()
   if not trial:
      FH = os.popen( cmd )
      output = [ line.rstrip() for line in FH.readlines() ]
      FH.close()
      return output
   return []

def SystemWrapper( cmd, verbose, trial = False ):
   dt = 0
   if verbose:
      PrintTimeStamp()
      print '... os.system(', cmd, ')'
      sys.stdout.flush()
   if trial:
      status = 0
   else:
      t0 = time.time()
      status = os.system( cmd )
      t1 = time.time()
      dt = t1 - t0
   assert status == 0
   if verbose:
      print '...', dt, 'seconds elapsed during system call.'
   return dt

def RemoveWrapper( fn, verbose, trial = False ):
	if verbose:
		PrintTimeStamp()
		print '... os.remove(', fn, ')'
		sys.stdout.flush()
	if not trial:
		os.remove( fn )

################################################################################
################################################################################
def ReadFileLines( file ):
   # opens a file
   # returns the contents as a list of lines in the file
   # trailing new-line chars are removed
   FH = open( file, 'r' )
   lines = [ line.rstrip() for line in FH.readlines() ]
   FH.close()
   return lines

################################################################################
################################################################################
def ExtractRawNameAndPath( fn ):
	( filePath, fn ) = os.path.split( fn )
	rawName =  '.'.join( fn.split('.')[:-1] )
	return ( filePath, rawName )

def GenTmpFileName( mtxFile ):
	( filePath, rawName ) = ExtractRawNameAndPath( mtxFile )
	return os.path.join( filePath, rawName + '.tmp.mtx' )

def GenDagFileName( mtxName, fmt, deg, numParts ):
   numParts = int( numParts )
   deg      = int( deg      )
   fext = '.%s.deg%02d.parts%02d.dag' % ( fmt.lower(), deg, numParts )
   ( filePath, rawName ) = ExtractRawNameAndPath( mtxFile )
   return os.path.join( filePath, rawName + fext )

def GenDupFileName( dagFile ):
	( filePath, rawName ) = ExtractRawNameAndPath( dagFile )
	toks = rawName.split( '.' )
	toks.pop()
	rawName = '.'.join( toks )
	return os.path.join( filePath, rawName + '.dup.mtx' )

################################################################################
################################################################################
def HeadNLines( fn, n ):
   cmd = ' '.join([ 'head', '-%d' % n, fn ])
   return SystemWrapperPipelined( cmd, False, False )

def VerifyMtxFileLine0( line0 ):
	toks  = line0.split()
	# ['MatrixMarket', 'matrix', 'coordinate', 'real', 'general']
	assert len( toks ) == 5
	assert toks[0] == 'MatrixMarket'
	assert toks[1] == 'matrix'
	assert toks[2] == 'coordinate'
	assert toks[3] == 'real'
	assert toks[4] == 'general'

def VerifyNonSymmRealMtx( mtxFile ):
	line0 = HeadNLines( mtxFile, 1 )[0].rstrip().lstrip( '%' )
	VerifyMtxFileLine0( line0 )

def VerifyMtxFileLine0SymmOblivious( line0 ):
	toks  = line0.split()
	# ['MatrixMarket', 'matrix', 'coordinate', 'real', 'general']
	assert len( toks ) == 5
	assert toks[0] == 'MatrixMarket'
	assert toks[1] == 'matrix'
	assert toks[2] == 'coordinate'
	assert toks[3] == 'real'
	assert toks[4] == 'general' or toks[4] == 'symmetric'

def VerifyRealMtx( mtxFile ):
	line0 = HeadNLines( mtxFile, 1 )[0].rstrip().lstrip( '%' )
	VerifyMtxFileLine0SymmOblivious( line0 )

################################################################################
################################################################################
def GenSequence( mtxFile, deg, numParts, params ):
   dagFile = GenDagFileName( mtxFile, params.fmt, deg, numParts )
   dupFile = GenDupFileName( dagFile )
   
   dagDump = ' '.join([ params.hcdDumpBin, params.fmt, mtxFile, dagFile, deg, numParts ])
   dagRead = ' '.join([ params.dagReadBin, params.fmt, dagFile, dupFile, deg, numParts ])
   mtxDiff = ' '.join([ params.mtxDiffBin, dupFile, mtxFile ])
   
   try:
      SystemWrapper( dagDump, params.verbose, params.trial )
      SystemWrapper( dagRead, params.verbose, params.trial )
      SystemWrapper( mtxDiff, params.verbose, params.trial )
   except:
      print ''
      print dagDump
      print dagRead
      print mtxDiff
      print ''
      if params.cleanup and os.path.isfile( dagFile ):
         RemoveWrapper( dagFile, params.verbose, params.trial )
         pass
      raise
   
   if params.cleanup:
      RemoveWrapper( dupFile, params.verbose, params.trial )

def RMASequence( mtxFile, deg, numParts, params ):
   assert int( numParts ) == 1
   GenSequence( mtxFile, deg, numParts, params )

def QTSSequence( mtxFile, deg, numParts, params ):
   assert int( numParts ) == 1
   GenSequence( mtxFile, deg, numParts, params )

def NZDSequence( mtxFile, deg, numParts, params ):
   assert int( numParts ) == 1
   GenSequence( mtxFile, deg, numParts, params )

def CSRSequence( mtxFile, deg, numParts, params ):
   GenSequence( mtxFile, deg, numParts, params )

def COOSequence( mtxFile, deg, numParts, params ):
   GenSequence( mtxFile, deg, numParts, params )

################################################################################
################################################################################
if len( sys.argv ) < 5:
   ( progPath, progName ) = os.path.split( sys.argv[0] )
   print 'Usage:', progName, '[RMA|QTS|NZD|CSR|COO] <mtxFile> <deg> <numParts>'
   sys.exit( -1 )

params = Params()

params.fmt = sys.argv[1]
mtxFile    = sys.argv[2]
deg        = str( int( sys.argv[3] ) )
numParts   = str( int( sys.argv[4] ) )

VerifyRealMtx( mtxFile )
if   params.fmt == 'RMA': RMASequence( mtxFile, deg, numParts, params )
elif params.fmt == 'QTS': QTSSequence( mtxFile, deg, numParts, params )
elif params.fmt == 'NZD': NZDSequence( mtxFile, deg, numParts, params )
elif params.fmt == 'CSR': CSRSequence( mtxFile, deg, numParts, params )
elif params.fmt == 'COO': COOSequence( mtxFile, deg, numParts, params )
else: assert False
